


from langgraph.graph import StateGraph, END
from typing import TypedDict, List, Annotated
import operator

from qwen_config import get_qwen_client

# 初始化通义千问客户端
client = get_qwen_client()



# 定义状态结构
class AgentState(TypedDict):
    # 消息列表，自动累积
    messages: Annotated[List[str], operator.add]
    current_response: str

def call_qwen_model(state: AgentState):
    """调用通义千问模型生成回复"""
    # global client # 全局变量在使用前初始化了就不需要再声明
    try:
        user_input = state['messages'][-1]

        response = client.chat.completions.create(
            model = "qwen-max",
            messages = [
                {"role":"system", "content":"你是一个有用的AI助手。"},
                {"role": "user", "content": user_input}
            ],
            temperature = 0.7
        )

        model_response = response.choices[0].message.content
        return {"current_response": model_response, "messages": [model_response]}
    
    except Exception as e:
        err = f"调用通义千问API出错：{str(e)}"
        return {"current_response": err, "messages": [err]}

def format_output(state: AgentState):
    """格式化输出"""
    response = state["current_response"]
    print(f"通义千问回复：{response}")
    return state

def create_workflow():
    """构建工作流"""
    workflow = StateGraph(AgentState)

    # 添加节点
    workflow.add_node("qwen_processor", call_qwen_model)
    workflow.add_node("output_formatter", format_output)

    # 设置入口点
    workflow.set_entry_point("qwen_processor")

    # 添加边
    workflow.add_edge("qwen_processor", "output_formatter")
    workflow.add_edge("output_formatter", END)

    return workflow.compile()

def main():
   
    graph = create_workflow()
    result = graph.invoke({"messages": ["请300字以介绍下你自己。"]})

    print("最终结果：", result)



if __name__ == "__main__":
    main()
